'''
Methods to facilitate viewing coverage data in the UCSC genome browser

Created on Nov 30, 2009

@author: mkiyer
'''
__author__ = 'mkiyer'
__version__ = '0.0.1'

import logging
import collections
import os
import sys
from optparse import OptionParser
from veggie.db.sample.samplegroup import parse_samplegroups_xml
from veggie.db.sample import get_sampledb_instance
from matplotlib.colors import ColorConverter

def get_rnaseq_tracks_path():
    path = os.path.join(os.path.sep, 'exds', 'users', 'mkiyer', 'tracks', 'rnaseq')
    return path

def get_rnaseq_tracks_url():
    ip_address = 'http://141.214.4.250'
    return os.path.join(ip_address, 'users', 'mkiyer', 'tracks', 'rnaseq')

def get_rnaseq_track_subdirs():
    return ['.']

def get_chipseq_tracks_path():
    path = os.path.join(os.path.sep, 'exds', 'users', 'mkiyer', 'tracks', 'chipseq')
    return path

def get_chipseq_tracks_url():
    ip_address = 'http://141.214.4.250'
    return os.path.join(ip_address, 'exds', 'users', 'mkiyer', 'tracks', 'chipseq')

def get_chipseq_track_subdirs():
    return ['.']

def get_default_track_params():
    return {'alwaysZero': 'on'}
    return {'autoScale': 'on',
            'alwaysZero': 'on'}

def find_file_subpath(filename, root_path):
    for subdir in get_rnaseq_track_subdirs():
        if os.path.exists(os.path.join(root_path, subdir, filename)):
            return os.path.join(subdir, filename)
    return None
    
def fetch_tracks_from_name(sample_name, sdbi, tracks_path, qc_filter=True):
    logger = logging.getLogger(__name__)    
    # get the libraries associated with this sample
    libraries = sdbi.get_libraries_by_sample_name(sample_name, best=True)
    for library in libraries:
        if library.qc_status == False:
            logger.warning('library %s: marked QC-FAIL' % library.id)
            if qc_filter:
                logger.warning('library %s: excluded from query' % library.id)
                continue
        # check for bigwig/bigbed files
        bwtrack = ('bigWig', 'coverage', library.id + '.bw')
        bbtrack = ('bigBed', 'junctions', library.id + '_junctions.bb')
        for track_info in [bwtrack, bbtrack]:
            track_type, track_desc, filename = track_info 
            filepath = find_file_subpath(filename, tracks_path)
            if filepath is None:
                logger.error('library %s file %s not found in database' % (library.id, filename))
            else:
                logger.error('library %s: found file %s' % (library.id, filename))
                yield (library.id, track_type, track_desc, filepath)

def fetch_tracks_from_xml(sample_xmlfile, sdbi, tracks_path, qc_filter=True):
    # load samples and construct database queries
    current_priority = 1
    tracks = collections.defaultdict(lambda: [])
    for sgroup in parse_samplegroups_xml(sample_xmlfile):
        for sample_name in sgroup.samples:
            # get the libraries associated with this sample            
            for library_id, track_type, track_desc, track_path in fetch_tracks_from_name(sample_name, sdbi, tracks_path, qc_filter):
                tracks[sample_name].append((library_id, sgroup.color, current_priority, track_type, track_desc, track_path))
                current_priority += 1
    return tracks

def make_track_header(name, type, url, desc=None, priority=None, color=None, **kwargs):
    '''
    http://genome.ucsc.edu/goldenPath/help/bedgraph.html

    track type=bedGraph name=track_label description=center_label
        visibility=display_mode color=r,g,b altColor=r,g,b
        priority=priority autoScale=on|off alwaysZero=on|off
        gridDefault=on|off maxHeightPixels=max:default:min
        graphType=bar|points viewLimits=lower:upper
        yLineMark=real-value yLineOnOff=on|off
        windowingFunction=maximum|mean|minimum smoothingWindow=off|2-16
  
    chromA  chromStartA  chromEndA  dataValueA
    chromB  chromStartB  chromEndB  dataValueB
    '''    
    if desc is None:
        desc = name
    track_header = ['track',
                    'type=%s' % type,
                    'name="%s"' % name,
                    'description="%s"' % desc,
                    'bigDataUrl=%s' % (url),
                    'visibility=full']
    for k,v in kwargs.iteritems():
        track_header.append('%s=%s' % (k,v))
    if color is not None:
        track_header.append('color=%s' % color)
    if priority is not None:
        track_header.append('priority=%d' % priority)
    return ' '.join(track_header)

def make_tracks(tracks, tracks_baseurl, outfhd, **kwargs):
    logger = logging.getLogger(__name__)    
    logger.info("view_ucsc_tracks.py version %s" % (__version__))
    logger.info("----------------------------------")
    # make tracks
    color_converter = ColorConverter()
    for sample_name, sample_tracks in tracks.iteritems():        
        for track_info in sample_tracks:
            library_id, library_color, priority, track_type, track_desc, track_path= track_info            
            # setup track color
            track_rgb = color_converter.to_rgb(library_color)
            r, g, b = [int(255 * x) for x in track_rgb]
            track_color_str = ','.join(map(str, (r, g, b)))
            # make track header
            track_name = '_'.join([sample_name, library_id, track_desc])
            track_url = os.path.join(tracks_baseurl, track_path)   
            track_header = make_track_header(track_name, 
                                             '%s' % (track_type), 
                                             track_url, 
                                             desc=track_name,
                                             priority=priority,
                                             color=track_color_str,
                                             **kwargs)
            # output track headers
            outfhd.write(track_header + '\n')

def fetch_chipseq_tracks(tracks_path, tracks_baseurl, outfhd, **kwargs):
    import glob    
    for subdir in get_chipseq_track_subdirs():
        for filename in glob.iglob(os.path.join(tracks_path, subdir, "*.bw")):
            track_name = os.path.basename(filename).split(os.path.extsep)[0]            
            track_type = "bigWig"
            track_url = os.path.join(tracks_baseurl, subdir, os.path.basename(filename))            
            track_header = make_track_header(track_name, 
                                             track_type, 
                                             track_url, 
                                             desc=track_name,
                                             **kwargs)
            # output track headers
            outfhd.write(track_header + '\n')

if __name__ == '__main__':
    # setup logging
    logging.basicConfig(format='%(asctime)s %(levelname)s %(message)s',
                        level=logging.DEBUG)
    logger = logging.getLogger(__name__)

    # parse command line
    optionparser = OptionParser("usage: %prog [-i <chrN:start-end> -s <samples.xml>] [sample1, sample2, ...]")
    optionparser.add_option("--rnaseq-tracks", dest="rnaseq_tracks", help="path to ucsc tracks [default: %default]",
                            default=get_rnaseq_tracks_path())
    optionparser.add_option("--rnaseq-url", dest="rnaseq_url", 
                            help="url associated with tracks [default: %default]",
                            default=get_rnaseq_tracks_url())    
    optionparser.add_option("-s", "--samples", dest="samples", default=None,
                            help="sample group XML file")
    optionparser.add_option("-i", "--interval", dest="interval", default=None,
                            help="initial browser position (chrN:XXXX:YYYY)")
    optionparser.add_option("-o", "--output", dest="outfile",
                            default=None,
                            help="output file [default: %default]")
    optionparser.add_option("--chipseq", dest="chipseq", action="store_true", default=False)
    (options, args) = optionparser.parse_args()

    # process command line args
    rnaseq_path = options.rnaseq_tracks
    rnaseq_baseurl = options.rnaseq_url

    # TODO: add cmd line args for chipseq
    chipseq_path = get_chipseq_tracks_path()
    chipseq_baseurl = get_chipseq_tracks_url()
    
    if options.outfile == None:
        outfhd = sys.stdout
    else:
        outfhd = open(options.outfile, 'w')
    # load the current sampledb snapshot
    sdbi = get_sampledb_instance()
    # get the tracks
    tracks = collections.defaultdict(lambda: [])
    if options.samples is not None:
        # fetch samples from xml
        tracks.update(fetch_tracks_from_xml(options.samples, sdbi, rnaseq_path, qc_filter=True))
    # add any command-line samples
    current_priority = 100
    for arg in args:
        for library_id, track_type, track_desc, track_path in fetch_tracks_from_name(arg, sdbi, rnaseq_path, qc_filter=True):
            tracks[arg].append((library_id, 'black', current_priority, track_type, track_desc, track_path))
            current_priority += 1

    # get the browser position if specified on command line
    if options.interval is not None:
        # output browser position
        outfhd.write('browser position %s\n' % options.interval)
        #chrom, interval = interval_str.split(':')
        #start, end = map(int, interval.split('-'))        
    # construct track headers
    make_tracks(tracks, rnaseq_baseurl, outfhd, **get_default_track_params())
    
    # TODO: integrate chipseq with rnaseq tracks, but for now just output them (hack)
    if options.chipseq:
        track_params = {'alwaysZero': 'on',
                        'autoScale': 'off',
                        'viewLimits': '0:50',
                        'maxHeightPixels': '64:64:11'}
        fetch_chipseq_tracks(chipseq_path, chipseq_baseurl, outfhd, **track_params)

#http://genome.ucsc.edu/cgi-bin/hgTracks?org=organism_name&position=chr_position&hgt.customText=URL. 
#http://genome.ucsc.edu/cgi-bin/hgTracks?db=hg18&position=chr_position&hgt.customText=URL. 
#chr22:15916196-31832390
#http://genome.ucsc.edu/cgi-bin/hgTracks?org=human&db=hg18&hgt.customText=http://141.214.6.100//coverage_map_cache/tmpI6eaOG.bed.gz
